package com.zyb.lmbackend.llm;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.base.Stopwatch;
import com.zyb.lmbackend.config.AiProperties;
import com.zyb.lmbackend.constants.LlmPrompts;
import com.zyb.lmbackend.entity.LlmCallAudit;
import com.zyb.lmbackend.llm.model.CategoryDef;
import com.zyb.lmbackend.llm.model.ClsResult;
import com.zyb.lmbackend.service.LlmCallAuditService;
import com.zyb.lmbackend.util.InvalidLlmOutputException;
import com.zyb.lmbackend.util.JsonLinesOrArrayParser;
import com.zyb.lmbackend.util.JsonSchemaValidatorUtil;
import io.micrometer.core.instrument.MeterRegistry;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestTemplate;

/**
 * 基于 DashScope 应用(App) 完成接口的 LLM 客户端实现。
 *
 * 通过 /apps/{appId}/completion 调用，由平台应用统一注入系统提示与工作流。
 */
@RequiredArgsConstructor
@Slf4j
public class BailianLlmClient implements LlmClient {

    private final AiProperties aiProperties;
    private final ObjectMapper objectMapper = new ObjectMapper();
    private final LlmCallAuditService auditService;
    private final MeterRegistry meterRegistry;
    private final RestTemplate http = new RestTemplate();

    private static final String APPS_COMPLETION_ENDPOINT = "https://dashscope.aliyuncs.com/api/v1/apps/%s/completion";

    @Override
    /** 调用 App Completion 进行类目归纳。 */
    public List<CategoryDef> buildTaxonomy(List<String> samples) {
        var cfg = aiProperties.taxonomyCfg();
        String inPayload = toJsonSafe(new HashMap<String, Object>() {{ put("samples", samples); }});
        String responseText = callAgent(cfg, LlmPrompts.TAXONOMY_SYSTEM_PROMPT, inPayload);

        Stopwatch sw = Stopwatch.createStarted();
        Integer status = 200;
        try {
            List<CategoryDef> result = objectMapper.readValue(responseText, new TypeReference<List<CategoryDef>>() {});
            return result;
        } catch (Exception e) {
            log.error("Parse taxonomy output failed: {}", truncate64k(responseText));
            meterRegistry.counter("llm.errors").increment();
            throw new InvalidLlmOutputException("无法解析类目归纳结果", e);
        } finally {
            long latency = sw.stop().elapsed(TimeUnit.MILLISECONDS);
            meterRegistry.counter("llm.calls").increment();
            meterRegistry.timer("llm.latency").record(latency, TimeUnit.MILLISECONDS);
            saveAudit(null, null, "IN", inPayload, null, null, null);
            saveAudit(null, null, "OUT", truncate64k(responseText), null, (int) latency, status);
        }
    }

    @Override
    /** 调用 App Completion 进行批量分类。 */
    public List<ClsResult> classifyBatch(List<String> details, List<String> allowedCodes) {
        var cfg = aiProperties.classifyCfg();
        String sys = LlmPrompts.CLASSIFY_SYSTEM_PROMPT + "\n允许的code列表�?" + String.join(",", allowedCodes);
        String inPayload = toJsonSafe(details);
        String raw = callAgent(cfg, sys, inPayload);

        Stopwatch sw = Stopwatch.createStarted();
        Integer status = 200;
        try {
            JsonLinesOrArrayParser parser = new JsonLinesOrArrayParser(objectMapper);
            List<JsonNode> nodes = parser.parseToNodes(raw);
            new JsonSchemaValidatorUtil(objectMapper).validateEach(nodes, "schema/classify_item.schema.json");
            List<ClsResult> list = new ArrayList<>();
            for (JsonNode n : nodes) {
                int index = n.path("index").asInt();
                String code = n.path("category_code").asText();
                double conf = n.path("confidence").asDouble();
                String reason = n.path("reason").asText();
                list.add(new ClsResult(index, code, conf, reason));
            }
            return list;
        } catch (Exception e) {
            log.error("Parse classify output failed: {}", truncate64k(raw));
            meterRegistry.counter("llm.errors").increment();
            throw new InvalidLlmOutputException("无法解析分类结果", e);
        } finally {
            long latency = sw.stop().elapsed(TimeUnit.MILLISECONDS);
            meterRegistry.counter("llm.calls").increment();
            meterRegistry.timer("llm.latency").record(latency, TimeUnit.MILLISECONDS);
            saveAudit(null, null, "IN", new Object() { public final List<String> items = details; public final List<String> codes = allowedCodes; }.toString(), null, null, null);
            saveAudit(null, null, "OUT", truncate64k(raw), null, (int) latency, status);
        }
    }

    /**
     * 调用 App Completion 接口并尽力提取 message.content。
     */
    private String callAgent(AiProperties.AgentCfg cfg, String systemPrompt, String userContent) {
        if (cfg == null) throw new IllegalStateException("DashScope Agent 未配�?");
        String url = String.format(APPS_COMPLETION_ENDPOINT, cfg.getAppId());

        Map<String, Object> sysMsg = Map.of("role", "system", "content", systemPrompt);
        Map<String, Object> userMsg = Map.of("role", "user", "content", userContent);
        Map<String, Object> body = new HashMap<>();
        body.put("input", Map.of("messages", List.of(sysMsg, userMsg)));
        Map<String, Object> params = new HashMap<>();
        if (cfg.getTemperature() != null) params.put("temperature", cfg.getTemperature());
        params.put("incremental_output", false);
        body.put("parameters", params);

        String inJson;
        try { inJson = objectMapper.writeValueAsString(body); } catch (JsonProcessingException e) { inJson = String.valueOf(body); }

        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.setBearerAuth(cfg.getApiKey());
        HttpEntity<String> entity = new HttpEntity<>(inJson, headers);

        ResponseEntity<String> resp = http.postForEntity(url, entity, String.class);
        String bodyText = resp.getBody() == null ? "" : resp.getBody();
        try {
            JsonNode root = objectMapper.readTree(bodyText);
            JsonNode content = null;
            if (root.has("output")) {
                JsonNode out = root.get("output");
                if (out.has("choices")) {
                    JsonNode arr = out.get("choices");
                    if (arr.isArray() && arr.size() > 0) {
                        JsonNode ch0 = arr.get(0);
                        if (ch0.has("message")) {
                            content = ch0.get("message").get("content");
                        }
                    }
                }
                if (content == null && out.has("text")) content = out.get("text");
                if (content == null && out.has("message")) content = out.get("message").get("content");
            }
            if (content == null && root.has("choices")) {
                JsonNode arr = root.get("choices");
                if (arr.isArray() && arr.size() > 0) {
                    JsonNode ch0 = arr.get(0);
                    if (ch0.has("message")) content = ch0.get("message").get("content");
                }
            }
            if (content != null && !content.isNull()) {
                return content.asText();
            }
        } catch (Exception ignore) {
        }
        return bodyText;
    }

    private void saveAudit(Long taskId, Integer batchNo, String direction, String payload,
            Integer tokens, Integer latencyMs, Integer httpStatus) {
        LlmCallAudit audit = new LlmCallAudit();
        audit.setTaskId(taskId);
        audit.setBatchNo(batchNo);
        audit.setDirection(direction);
        audit.setPayload(truncate64k(payload));
        audit.setTokens(tokens);
        audit.setLatencyMs(latencyMs);
        audit.setHttpStatus(httpStatus);
        auditService.save(audit);
    }

    private String toJsonSafe(Object obj) {
        try {
            return objectMapper.writeValueAsString(obj);
        } catch (JsonProcessingException e) {
            return String.valueOf(obj);
        }
    }

    private String truncate64k(String s) {
        if (s == null) return null;
        byte[] bytes = s.getBytes(StandardCharsets.UTF_8);
        int limit = 64 * 1024; // 64KB
        if (bytes.length <= limit) return s;
        int end = limit;
        while (end > 0 && (bytes[end - 1] & 0xC0) == 0x80) {
            end--;
        }
        return new String(bytes, 0, end, StandardCharsets.UTF_8);
    }
}

